package rpc

import (
	"bestgo/pkg/codec"
	"compress/flate"
	"net"
	"net/rpc"
	"sync"
	"time"
)

type SmartConnectionPool struct {
	pool  chan *smartClient
	addr  string
	stats struct {
		totalLatency time.Duration
		requestCount int
	}
	statLock sync.Mutex
}

type smartClient struct {
	*rpc.Client
	latency time.Duration
}

func NewSmartConnectionPool(addr string, size int) *SmartConnectionPool {
	p := &SmartConnectionPool{
		pool: make(chan *smartClient, size),
		addr: addr,
	}

	for i := 0; i < size; i++ {
		conn, _ := net.Dial("tcp", addr)
		client := &smartClient{
			Client: rpc.NewClientWithCodec(
				codec.NewCompressedGobCodec(conn, flate.DefaultCompression)),
		}
		p.pool <- client
	}
	return p
}

func (p *SmartConnectionPool) Call(method string, args any, reply any) error {
	start := time.Now()
	latency := time.Since(start)

	client := <-p.pool
	defer func() {
		p.pool <- client

		// 更新客户端延迟数据
		client.latency = latency

		// 更新全局统计
		p.statLock.Lock()
		p.stats.totalLatency += latency
		p.stats.requestCount++
		p.statLock.Unlock()
	}()

	err := client.Call(method, args, reply)
	if err != nil {
		return err
	}

	return err
}

func (p *SmartConnectionPool) AvgLatency() time.Duration {
	p.statLock.Lock()
	defer p.statLock.Unlock()

	if p.stats.requestCount == 0 {
		return 0
	}
	return p.stats.totalLatency / time.Duration(p.stats.requestCount)
}

func (c *SmartConnectionPool) Metrics() map[string]interface{} {
	return map[string]interface{}{
		"avgLagtency": c.AvgLatency(),
	}
}
